import numpy as np
import matplotlib.pyplot as plt


def model(x, theta):
    return x.dot(theta)


def cost(h, y):
    return 0.5 * np.mean((h - y) ** 2)


def grad(x, y, iter0=5000, alpha=0.01):
    m, n = x.shape
    theta = np.zeros(n)
    J = np.zeros(iter0)
    for i in range(iter0):
        h = model(x, theta)
        J[i] = cost(h, y)
        dt = 1 / m * x.T.dot(h - y)
        theta -= alpha * dt
    return theta, J, h


def score(h, y):
    u = np.sum((h-y) ** 2)
    mu = np.mean(y)
    v = np.sum((y-mu) ** 2)
    return 1 - u/v


if __name__ == '__main__':
    data = np.loadtxt('../data/ex1data2.txt', delimiter=',')

    x = data[:, :-1]
    y = data[:, -1]
    m = len(x)

    # 归一化
    min_x=np.min(x,axis=0)
    max_x=np.max(x,axis=0)
    x=(x-min_x)/(max_x-min_x)

    # #标准化
    # miu = np.mean(x, axis=0)
    # sigma = np.std(x, axis=0)
    # x = (x - miu) / sigma

    # 拼接
    X = np.c_[np.ones(m), x]

    # 洗牌
    np.random.seed(666)
    a = np.random.permutation(m)
    X = X[a]
    y = y[a]

    # 切分 训练集、测试集
    num = int(0.7 * m)
    train_X, test_X = np.split(X, [num, ])
    train_y, test_y = np.split(y, [num, ])

    # 训练模型
    theta, J, train_h = grad(train_X, train_y)
    plt.figure(figsize=[12, 5])
    spr = 1
    spc = 3
    plt.subplot(spr, spc, 1)
    plt.plot(J)
    print('THETA', theta)
    print('Score train', score(train_h, train_y))
    test_h = model(test_X, theta)
    print('Score test', score(test_h, test_y))

    plt.subplot(spr, spc, 2)
    plt.scatter(train_y, train_y)
    plt.scatter(train_y, train_h)

    plt.subplot(spr, spc, 3)
    plt.scatter(test_y, test_y)
    plt.scatter(test_y, test_h)

    plt.show()
